{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#!/usr/bin/env python\n",
    "# coding: utf-8\n",
    "\n",
    "import os\n",
    "import math\n",
    "import time\n",
    "import json\n",
    "import random\n",
    "\n",
    "from collections import OrderedDict\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from data.data_iterator import TextIterator\n",
    "from data.data_iterator import BiTextIterator\n",
    "\n",
    "import data.data_utils as data_utils\n",
    "from data.data_utils import prepare_batch\n",
    "from data.data_utils import prepare_train_batch\n",
    "\n",
    "from seq2seq_model import Seq2SeqModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Data loading parameters\n",
    "tf.app.flags.DEFINE_string('source_vocabulary', 'data/europarl-v7.1.4M.de.json', 'Path to source vocabulary')\n",
    "tf.app.flags.DEFINE_string('target_vocabulary', 'data/europarl-v7.1.4M.fr.json', 'Path to target vocabulary')\n",
    "tf.app.flags.DEFINE_string('source_train_data', 'data/europarl-v7.1.4M.de', 'Path to source training data')\n",
    "tf.app.flags.DEFINE_string('target_train_data', 'data/europarl-v7.1.4M.fr', 'Path to target training data')\n",
    "tf.app.flags.DEFINE_string('source_valid_data', 'data/newstest2012.bpe.de', 'Path to source validation data')\n",
    "tf.app.flags.DEFINE_string('target_valid_data', 'data/newstest2012.bpe.fr', 'Path to target validation data')\n",
    "\n",
    "# Network parameters\n",
    "tf.app.flags.DEFINE_string('cell_type', 'lstm', 'RNN cell for encoder and decoder, default: lstm')\n",
    "tf.app.flags.DEFINE_string('attention_type', 'bahdanau', 'Attention mechanism: (bahdanau, luong), default: bahdanau')\n",
    "tf.app.flags.DEFINE_integer('hidden_units', 1024, 'Number of hidden units in each layer')\n",
    "tf.app.flags.DEFINE_integer('depth', 2, 'Number of layers in each encoder and decoder')\n",
    "tf.app.flags.DEFINE_integer('embedding_size', 500, 'Embedding dimensions of encoder and decoder inputs')\n",
    "tf.app.flags.DEFINE_integer('num_encoder_symbols', 30000, 'Source vocabulary size')\n",
    "tf.app.flags.DEFINE_integer('num_decoder_symbols', 30000, 'Target vocabulary size')\n",
    "\n",
    "tf.app.flags.DEFINE_boolean('use_residual', True, 'Use residual connection between layers')\n",
    "tf.app.flags.DEFINE_boolean('attn_input_feeding', False, 'Use input feeding method in attentional decoder')\n",
    "tf.app.flags.DEFINE_boolean('use_dropout', True, 'Use dropout in each rnn cell')\n",
    "tf.app.flags.DEFINE_float('dropout_rate', 0.3, 'Dropout probability for input/output/state units (0.0: no dropout)')\n",
    "\n",
    "# Training parameters\n",
    "tf.app.flags.DEFINE_float('learning_rate', 0.0002, 'Learning rate')\n",
    "tf.app.flags.DEFINE_float('max_gradient_norm', 1.0, 'Clip gradients to this norm')\n",
    "tf.app.flags.DEFINE_integer('batch_size', 128, 'Batch size')\n",
    "tf.app.flags.DEFINE_integer('max_epochs', 10, 'Maximum # of training epochs')\n",
    "tf.app.flags.DEFINE_integer('max_load_batches', 20, 'Maximum # of batches to load at one time')\n",
    "tf.app.flags.DEFINE_integer('max_seq_length', 50, 'Maximum sequence length')\n",
    "tf.app.flags.DEFINE_integer('display_freq', 100, 'Display training status every this iteration')\n",
    "tf.app.flags.DEFINE_integer('save_freq', 11500, 'Save model checkpoint every this iteration')\n",
    "tf.app.flags.DEFINE_integer('valid_freq', 1150000, 'Evaluate model every this iteration: valid_data needed')\n",
    "tf.app.flags.DEFINE_string('optimizer', 'adam', 'Optimizer for training: (adadelta, adam, rmsprop)')\n",
    "tf.app.flags.DEFINE_string('model_dir', 'model/', 'Path to save model checkpoints')\n",
    "tf.app.flags.DEFINE_string('summary_dir', 'model/summary', 'Path to save model summary')\n",
    "tf.app.flags.DEFINE_string('model_name', 'translate.ckpt', 'File name used for model checkpoints')\n",
    "tf.app.flags.DEFINE_boolean('shuffle_each_epoch', True, 'Shuffle training dataset for each epoch')\n",
    "tf.app.flags.DEFINE_boolean('sort_by_length', True, 'Sort pre-fetched minibatches by their target sequence lengths')\n",
    "tf.app.flags.DEFINE_boolean('use_fp16', False, 'Use half precision float16 instead of float32 as dtype')\n",
    "\n",
    "# Runtime parameters\n",
    "tf.app.flags.DEFINE_boolean('allow_soft_placement', True, 'Allow device soft placement')\n",
    "tf.app.flags.DEFINE_boolean('log_device_placement', False, 'Log placement of ops on devices')\n",
    "\n",
    "FLAGS = tf.app.flags.FLAGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def create_model(session, FLAGS):\n",
    "\n",
    "    config = OrderedDict(sorted(FLAGS.__flags.items()))\n",
    "    model = Seq2SeqModel(config, 'train')\n",
    "\n",
    "    ckpt = tf.train.get_checkpoint_state(FLAGS.model_dir)\n",
    "    if ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path):\n",
    "        print 'Reloading model parameters..'\n",
    "        model.restore(session, ckpt.model_checkpoint_path)\n",
    "        \n",
    "    else:\n",
    "        if not os.path.exists(FLAGS.model_dir):\n",
    "            os.makedirs(FLAGS.model_dir)\n",
    "        print 'Created new model parameters..'\n",
    "        session.run(tf.global_variables_initializer())\n",
    "   \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train():\n",
    "    # Load parallel data to train\n",
    "    print 'Loading training data..'\n",
    "    train_set = BiTextIterator(source=FLAGS.source_train_data,\n",
    "                               target=FLAGS.target_train_data,\n",
    "                               source_dict=FLAGS.source_vocabulary,\n",
    "                               target_dict=FLAGS.target_vocabulary,\n",
    "                               batch_size=FLAGS.batch_size,\n",
    "                               maxlen=FLAGS.max_seq_length,\n",
    "                               n_words_source=FLAGS.num_encoder_symbols,\n",
    "                               n_words_target=FLAGS.num_decoder_symbols,\n",
    "                               shuffle_each_epoch=FLAGS.shuffle_each_epoch,\n",
    "                               sort_by_length=FLAGS.sort_by_length,\n",
    "                               maxibatch_size=FLAGS.max_load_batches)\n",
    "\n",
    "    if FLAGS.source_valid_data and FLAGS.target_valid_data:\n",
    "        print 'Loading validation data..'\n",
    "        valid_set = BiTextIterator(source=FLAGS.source_valid_data,\n",
    "                                   target=FLAGS.target_valid_data,\n",
    "                                   source_dict=FLAGS.source_vocabulary,\n",
    "                                   target_dict=FLAGS.target_vocabulary,\n",
    "                                   batch_size=FLAGS.batch_size,\n",
    "                                   maxlen=None,\n",
    "                                   n_words_source=FLAGS.num_encoder_symbols,\n",
    "                                   n_words_target=FLAGS.num_decoder_symbols)\n",
    "    else:\n",
    "        valid_set = None\n",
    "\n",
    "    # Initiate TF session\n",
    "    with tf.Session(config=tf.ConfigProto(allow_soft_placement=FLAGS.allow_soft_placement, \n",
    "        log_device_placement=FLAGS.log_device_placement, gpu_options=tf.GPUOptions(allow_growth=True))) as sess:\n",
    "\n",
    "        # Create a log writer object\n",
    "        log_writer = tf.summary.FileWriter(FLAGS.model_dir, graph=sess.graph)\n",
    "        \n",
    "        # Create a new model or reload existing checkpoint\n",
    "        model = create_model(sess, FLAGS)\n",
    "\n",
    "        step_time, loss = 0.0, 0.0\n",
    "        words_seen, sents_seen = 0, 0\n",
    "        start_time = time.time()\n",
    "\n",
    "        # Training loop\n",
    "        print 'Training..'\n",
    "        for epoch_idx in xrange(FLAGS.max_epochs):\n",
    "            if model.global_epoch_step.eval() >= FLAGS.max_epochs:\n",
    "                print 'Training is already complete.', \\\n",
    "                      'current epoch:{}, max epoch:{}'.format(model.global_epoch_step.eval(), FLAGS.max_epochs)\n",
    "                break\n",
    "\n",
    "            for source_seq, target_seq in train_set:    \n",
    "                # Get a batch from training parallel data\n",
    "                source, source_len, target, target_len = prepare_train_batch(source_seq, target_seq,\n",
    "                                                                             FLAGS.max_seq_length)\n",
    "                if source is None or target is None:\n",
    "                    print 'No samples under max_seq_length ', FLAGS.max_seq_length\n",
    "                    continue\n",
    "\n",
    "                # Execute a single training step\n",
    "                step_loss, summary = model.train(sess, encoder_inputs=source, encoder_inputs_length=source_len, \n",
    "                                                 decoder_inputs=target, decoder_inputs_length=target_len)\n",
    "\n",
    "                loss += float(step_loss) / FLAGS.display_freq\n",
    "                words_seen += float(np.sum(source_len+target_len))\n",
    "                sents_seen += float(source.shape[0]) # batch_size\n",
    "\n",
    "                if model.global_step.eval() % FLAGS.display_freq == 0:\n",
    "\n",
    "                    avg_perplexity = math.exp(float(loss)) if loss < 300 else float(\"inf\")\n",
    "\n",
    "                    time_elapsed = time.time() - start_time\n",
    "                    step_time = time_elapsed / FLAGS.display_freq\n",
    "\n",
    "                    words_per_sec = words_seen / time_elapsed\n",
    "                    sents_per_sec = sents_seen / time_elapsed\n",
    "\n",
    "                    print 'Epoch ', model.global_epoch_step.eval(), 'Step ', model.global_step.eval(), \\\n",
    "                          'Perplexity {0:.2f}'.format(avg_perplexity), 'Step-time ', step_time, \\\n",
    "                          '{0:.2f} sents/s'.format(sents_per_sec), '{0:.2f} words/s'.format(words_per_sec)\n",
    "\n",
    "                    loss = 0\n",
    "                    words_seen = 0\n",
    "                    sents_seen = 0\n",
    "                    start_time = time.time()\n",
    "                    \n",
    "                    # Record training summary for the current batch\n",
    "                    log_writer.add_summary(summary, model.global_step.eval())\n",
    "\n",
    "                # Execute a validation step\n",
    "                if valid_set and model.global_step.eval() % FLAGS.valid_freq == 0:\n",
    "                    print 'Validation step'\n",
    "                    valid_loss = 0.0\n",
    "                    valid_sents_seen = 0\n",
    "                    for source_seq, target_seq in valid_set:\n",
    "                        # Get a batch from validation parallel data\n",
    "                        source, source_len, target, target_len = prepare_train_batch(source_seq, target_seq)\n",
    "\n",
    "                        # Compute validation loss: average per word cross entropy loss\n",
    "                        step_loss = model.eval(sess, encoder_inputs=source, encoder_inputs_length=source_len,\n",
    "                                               decoder_inputs=target, decoder_inputs_length=target_len)\n",
    "                        batch_size = source.shape[0]\n",
    "\n",
    "                        valid_loss += step_loss * batch_size\n",
    "                        valid_sents_seen += batch_size\n",
    "                        print '  {} samples seen'.format(valid_sents_seen)\n",
    "\n",
    "                    valid_loss = valid_loss / valid_sents_seen\n",
    "                    print 'Valid perplexity: {0:.2f}'.format(math.exp(valid_loss))\n",
    "\n",
    "                # Save the model checkpoint\n",
    "                if model.global_step.eval() % FLAGS.save_freq == 0:\n",
    "                    print 'Saving the model..'\n",
    "                    checkpoint_path = os.path.join(FLAGS.model_dir, FLAGS.model_name)\n",
    "                    model.save(sess, checkpoint_path, global_step=model.global_step)\n",
    "                    json.dump(model.config,\n",
    "                              open('%s-%d.json' % (checkpoint_path, model.global_step.eval()), 'wb'),\n",
    "                              indent=2)\n",
    "\n",
    "            # Increase the epoch index of the model\n",
    "            model.global_epoch_step_op.eval()\n",
    "            print 'Epoch {0:} DONE'.format(model.global_epoch_step.eval())\n",
    "        \n",
    "        print 'Saving the last model..'\n",
    "        checkpoint_path = os.path.join(FLAGS.model_dir, FLAGS.model_name)\n",
    "        model.save(sess, checkpoint_path, global_step=model.global_step)\n",
    "        json.dump(model.config,\n",
    "                  open('%s-%d.json' % (checkpoint_path, model.global_step.eval()), 'wb'),\n",
    "                  indent=2)\n",
    "        \n",
    "    print 'Training Terminated'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def main(_):\n",
    "    train()\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    tf.app.run()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
